{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IRModule\n",
    "\n",
    "Apache TVM Unity 的核心抽象，即 IRModule。IRModule 包含整个 ML 模型，包括 计算图、张量程序和对外部库的潜在调用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/python/tvm_ffi/_optional_torch_c_dlpack.py:409: UserWarning: Failed to load torch c dlpack extension: Error building extension 'c_dlpack': [1/2] /media/pc/data/lxw/envs/anaconda3a/envs/py313/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=c_dlpack -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1018\\\" -I/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/include -I/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/3rdparty/dlpack/include -I/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/python/tvm_ffi/cython -I/media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include -I/media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include/torch/csrc/api/include -I/media/pc/data/lxw/envs/anaconda3a/envs/py313/include -isystem /media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include -isystem /media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include/torch/csrc/api/include -isystem /media/pc/data/lxw/envs/anaconda3a/envs/py313/include/python3.13 -fPIC -std=c++17 -O3 -DBUILD_WITH_CUDA -c /home/ai/.cache/torch_extensions/py313_cu128/c_dlpack/main.cpp -o main.o \n",
      "\u001b[31mFAILED: [code=1] \u001b[0mmain.o \n",
      "/media/pc/data/lxw/envs/anaconda3a/envs/py313/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=c_dlpack -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1018\\\" -I/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/include -I/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/3rdparty/dlpack/include -I/media/pc/data/lxw/ai/tvm/3rdparty/tvm-ffi/python/tvm_ffi/cython -I/media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include -I/media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include/torch/csrc/api/include -I/media/pc/data/lxw/envs/anaconda3a/envs/py313/include -isystem /media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include -isystem /media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include/torch/csrc/api/include -isystem /media/pc/data/lxw/envs/anaconda3a/envs/py313/include/python3.13 -fPIC -std=c++17 -O3 -DBUILD_WITH_CUDA -c /home/ai/.cache/torch_extensions/py313_cu128/c_dlpack/main.cpp -o main.o \n",
      "In file included from /home/ai/.cache/torch_extensions/py313_cu128/c_dlpack/main.cpp:8:\n",
      "/media/pc/data/lxw/envs/anaconda3a/envs/py313/lib/python3.13/site-packages/torch/include/c10/cuda/CUDAStream.h:3:10: fatal error: cuda_runtime_api.h: No such file or directory\n",
      "    3 | #include <cuda_runtime_api.h>\n",
      "      |          ^~~~~~~~~~~~~~~~~~~~\n",
      "compilation terminated.\n",
      "ninja: build stopped: subcommand failed.\n",
      ",EnvTensorAllocator will not be enabled.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['PATH'] += ':/usr/local/cuda/bin' # 保证 nvcc 可以被找到\n",
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 创建 IRModule\n",
    "\n",
    "IRModules 可以通过多种方式进行初始化。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 从现有前端模型导入\n",
    "初始化 IRModule 的最常见方法是从现有模型导入。Apache TVM Unity 支持从一系列框架导入，如 PyTorch 和 ONNX。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.export import export\n",
    "from tvm.relax.frontend.torch import from_exported_program"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(p_fc1_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(x, lv, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "            lv2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(lv1, p_fc1_bias)\n",
       "            lv3: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(lv2)\n",
       "            lv4: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(p_fc2_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            lv5: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(lv3, lv4, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "            lv6: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(lv5, p_fc2_bias)\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> lv6\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create a dummy model\n",
    "class TorchModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TorchModel, self).__init__()\n",
    "        self.fc1 = nn.Linear(784, 256)\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(256, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.relu1(x)\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# Give an example argument to torch.export\n",
    "example_args = (torch.randn(1, 784, dtype=torch.float32),)\n",
    "\n",
    "# Convert the model to IRModule\n",
    "with torch.no_grad():\n",
    "    exported_program = export(TorchModel().eval(), example_args)\n",
    "    mod_from_torch = from_exported_program(\n",
    "        exported_program, keep_params_as_input=True, unwrap_unit_return_tuple=True\n",
    "    )\n",
    "\n",
    "mod_from_torch, params_from_torch = relax.frontend.detach_params(mod_from_torch)\n",
    "# Print the IRModule\n",
    "mod_from_torch.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 使用 Relax NN 模块编写\n",
    "\n",
    "Apache TVM Unity还提供了一系列类似 PyTorch 的 API，帮助用户直接编写 IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">forward</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            permute_dims: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            matmul: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(x, permute_dims, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            add: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(matmul, fc1_bias)\n",
       "            relu: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(add)\n",
       "            permute_dims1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            matmul1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(relu, permute_dims1, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            add1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(matmul1, fc2_bias)\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> add1\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tvm.relax.frontend import nn\n",
    "\n",
    "\n",
    "class RelaxModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(RelaxModel, self).__init__()\n",
    "        self.fc1 = nn.Linear(784, 256)\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(256, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.relu1(x)\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "mod_from_relax, params_from_relax = RelaxModel().export_tvm(\n",
    "    {\"forward\": {\"x\": nn.spec.Tensor((1, 784), \"float32\")}}\n",
    ")\n",
    "mod_from_relax.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 通过 TVMScript 创建\n",
    "TVMScript 是一种基于 Python 的 DSL，用于 IRModule。我们可以直接以 TVMScript 语法输出 IRModule，或者解析 TVMScript 以获取 IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc1_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), fc2_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            permute_dims: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(fc1_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            matmul: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(x, permute_dims, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            add: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(matmul, fc1_bias)\n",
       "            relu: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(add)\n",
       "            permute_dims1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>permute_dims(fc2_weight, axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>)\n",
       "            matmul1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>matmul(relu, permute_dims1, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;void&quot;</span>)\n",
       "            add1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(matmul1, fc2_bias)\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> add1\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "\n",
    "\n",
    "@I.ir_module\n",
    "class TVMScriptModule:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x: R.Tensor((1, 784), dtype=\"float32\"),\n",
    "        fc1_weight: R.Tensor((256, 784), dtype=\"float32\"),\n",
    "        fc1_bias: R.Tensor((256,), dtype=\"float32\"),\n",
    "        fc2_weight: R.Tensor((10, 256), dtype=\"float32\"),\n",
    "        fc2_bias: R.Tensor((10,), dtype=\"float32\"),\n",
    "    ) -> R.Tensor((1, 10), dtype=\"float32\"):\n",
    "        R.func_attr({\"num_input\": 1})\n",
    "        with R.dataflow():\n",
    "            permute_dims = R.permute_dims(fc1_weight, axes=None)\n",
    "            matmul = R.matmul(x, permute_dims, out_dtype=\"void\")\n",
    "            add = R.add(matmul, fc1_bias)\n",
    "            relu = R.nn.relu(add)\n",
    "            permute_dims1 = R.permute_dims(fc2_weight, axes=None)\n",
    "            matmul1 = R.matmul(relu, permute_dims1, out_dtype=\"void\")\n",
    "            add1 = R.add(matmul1, fc2_bias)\n",
    "            gv = add1\n",
    "            R.output(gv)\n",
    "        return gv\n",
    "\n",
    "\n",
    "mod_from_script = TVMScriptModule\n",
    "mod_from_script.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IRModule的属性\n",
    "\n",
    "IRModule 是一组函数的集合，通过 GlobalVars 索引。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I.GlobalVar(\"main\")]\n"
     ]
    }
   ],
   "source": [
    "mod = mod_from_torch\n",
    "print(mod.get_global_vars())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以通过使用 GlobalVars 或它们的名称来索引 IRModule 中的函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# from tvm.script import relax as R\n",
      "\n",
      "@R.function\n",
      "def main(x: R.Tensor((1, 784), dtype=\"float32\"), p_fc1_weight: R.Tensor((256, 784), dtype=\"float32\"), p_fc1_bias: R.Tensor((256,), dtype=\"float32\"), p_fc2_weight: R.Tensor((10, 256), dtype=\"float32\"), p_fc2_bias: R.Tensor((10,), dtype=\"float32\")) -> R.Tensor((1, 10), dtype=\"float32\"):\n",
      "    R.func_attr({\"num_input\": 1})\n",
      "    with R.dataflow():\n",
      "        lv: R.Tensor((784, 256), dtype=\"float32\") = R.permute_dims(p_fc1_weight, axes=None)\n",
      "        lv1: R.Tensor((1, 256), dtype=\"float32\") = R.matmul(x, lv, out_dtype=\"float32\")\n",
      "        lv2: R.Tensor((1, 256), dtype=\"float32\") = R.add(lv1, p_fc1_bias)\n",
      "        lv3: R.Tensor((1, 256), dtype=\"float32\") = R.nn.relu(lv2)\n",
      "        lv4: R.Tensor((256, 10), dtype=\"float32\") = R.permute_dims(p_fc2_weight, axes=None)\n",
      "        lv5: R.Tensor((1, 10), dtype=\"float32\") = R.matmul(lv3, lv4, out_dtype=\"float32\")\n",
      "        lv6: R.Tensor((1, 10), dtype=\"float32\") = R.add(lv5, p_fc2_bias)\n",
      "        gv: R.Tensor((1, 10), dtype=\"float32\") = lv6\n",
      "        R.output(gv)\n",
      "    return gv\n"
     ]
    }
   ],
   "source": [
    "# index by global var name\n",
    "print(mod[\"main\"])\n",
    "# index by global var, and checking they are the same function\n",
    "(gv,) = mod.get_global_vars()\n",
    "assert mod[gv] == mod[\"main\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IRModule 上的变换\n",
    "\n",
    "变换是 Apache TVM Unity 的重要组成部分。一个变换接受一个 IRModule 并输出另一个 IRModule。我们可以将一系列变换应用于一个 IRModule 以获得一个新的 IRModule。这是优化模型的常见方法。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "有关每个变换的详细信息，请参阅[变换 API 参考](https://tvm.apache.org/docs/reference/api/python/relax/transform.html#api-relax-transformation)。\n",
    "\n",
    "首先对 IRModule 应用 `LegalizeOps` 变换。此变换将 Relax 模块转换为混合阶段，同一个模块内包 Relax 和 TensorIR 函数。同时，Relax 算子将被转换为 `call_tir`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">add</span>(lv1: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_bias: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_add: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(lv1[v_ax0, v_ax1], p_fc1_bias[v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> lv1[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">+</span> p_fc1_bias[v_ax1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">add1</span>(lv5: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_bias: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_add: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(lv5[v_ax0, v_ax1], p_fc2_bias[v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> lv5[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">+</span> p_fc2_bias[v_ax1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">matmul</span>(x: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), lv: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), matmul: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
       "                v_i0, v_i1, v_k <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(x[v_i0, v_k], lv[v_k, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(matmul[v_i0, v_i1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>init():\n",
       "                    matmul[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "                matmul[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> matmul[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">+</span> x[v_i0, v_k] <span style=\"color: #A2F; font-weight: bold\">*</span> lv[v_k, v_i1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">matmul1</span>(lv3: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), lv4: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), matmul: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
       "                v_i0, v_i1, v_k <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(lv3[v_i0, v_k], lv4[v_k, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(matmul[v_i0, v_i1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>init():\n",
       "                    matmul[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "                matmul[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> matmul[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">+</span> lv3[v_i0, v_k] <span style=\"color: #A2F; font-weight: bold\">*</span> lv4[v_k, v_i1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">relu</span>(lv2: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), compute: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(lv2[v_i0, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(compute[v_i0, v_i1])\n",
       "                compute[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(lv2[v_i0, v_i1], T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">transpose</span>(p_fc1_weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(p_fc1_weight[v_ax1, v_ax0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
       "                T_transpose[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> p_fc1_weight[v_ax1, v_ax0]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">transpose1</span>(p_fc2_weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(p_fc2_weight[v_ax1, v_ax0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
       "                T_transpose[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> p_fc2_weight[v_ax1, v_ax0]\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            lv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>transpose, (p_fc1_weight,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv1 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>matmul, (x, lv), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv2 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>add, (lv1, p_fc1_bias), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv3 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>relu, (lv2,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv4 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>transpose1, (p_fc2_weight,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv5 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>matmul1, (lv3, lv4), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv6 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>add1, (lv5, p_fc2_bias), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> lv6\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod = mod_from_torch\n",
    "mod = relax.transform.LegalizeOps()(mod)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "变换后，模块内会有更多的函数。再次打印全局变量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I.GlobalVar(\"add\"), I.GlobalVar(\"add1\"), I.GlobalVar(\"main\"), I.GlobalVar(\"matmul\"), I.GlobalVar(\"matmul1\"), I.GlobalVar(\"relu\"), I.GlobalVar(\"transpose\"), I.GlobalVar(\"transpose1\")]\n"
     ]
    }
   ],
   "source": [
    "print(mod.get_global_vars())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Apache TVM Unity 为用户提供了一组默认的变换管道，以简化变换过程。然后我们可以将默认管道应用于模块。默认的零管道包含一些非常基础的变换，包括：\n",
    "\n",
    "- `LegalizeOps`：此变换将 Relax 算子转换为具有相应 TensorIR 函数的 `call_tir` 函数。在此变换之后，IRModule 将包含 Relax 函数和 TensorIR 函数。\n",
    "- `AnnotateTIROpPattern`：此变换注解 TensorIR 函数的模式，为后续的算子融合做准备。\n",
    "- `FoldConstant`：此 pass 执行常量折叠，优化涉及常量的运算。\n",
    "- `FuseOps` 和 `FuseTIR`：这两个传递基于上一步（`AnnotateTIROpPattern`）中注解的模式共同工作以融合算子。这些传递转换 Relax 函数和 TensorIR 函数。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "在这里，我们在流程中应用了两次 `LegalizeOps`。第二次是多余的，但无害。\n",
    "\n",
    "每个传递都可以在流程中重复，因为我们确保传递可以处理所有合法的 IRModule 输入。这种设计可以帮助用户构建他们自己的管道。   \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">fused_matmul1_add1</span>(lv3: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), lv4: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_bias: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_add_intermediate: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        matmul_intermediate <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>alloc_buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)))\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
       "                v_i0, v_i1, v_k <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(lv3[v_i0, v_k], lv4[v_k, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(matmul_intermediate[v_i0, v_i1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>init():\n",
       "                    matmul_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "                matmul_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> matmul_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">+</span> lv3[v_i0, v_k] <span style=\"color: #A2F; font-weight: bold\">*</span> lv4[v_k, v_i1]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(matmul_intermediate[v_ax0, v_ax1], p_fc2_bias[v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add_intermediate[v_ax0, v_ax1])\n",
       "                T_add_intermediate[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> matmul_intermediate[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">+</span> p_fc2_bias[v_ax1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">fused_matmul_add_relu</span>(x: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), lv: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_bias: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>),), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), compute_intermediate: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        matmul_intermediate <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>alloc_buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)))\n",
       "        T_add_intermediate <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>alloc_buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)))\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;matmul&quot;</span>):\n",
       "                v_i0, v_i1, v_k <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [i0, i1, k])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(x[v_i0, v_k], lv[v_k, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(matmul_intermediate[v_i0, v_i1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>init():\n",
       "                    matmul_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
       "                matmul_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> matmul_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">+</span> x[v_i0, v_k] <span style=\"color: #A2F; font-weight: bold\">*</span> lv[v_k, v_i1]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(matmul_intermediate[v_ax0, v_ax1], p_fc1_bias[v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add_intermediate[v_ax0, v_ax1])\n",
       "                T_add_intermediate[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> matmul_intermediate[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">+</span> p_fc1_bias[v_ax1]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(T_add_intermediate[v_i0, v_i1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(compute_intermediate[v_i0, v_i1])\n",
       "                compute_intermediate[v_i0, v_i1] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(T_add_intermediate[v_i0, v_i1], T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">transpose</span>(p_fc1_weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;op_pattern&quot;</span>: <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(p_fc1_weight[v_ax1, v_ax0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
       "                T_transpose[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> p_fc1_weight[v_ax1, v_ax0]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">transpose1</span>(p_fc2_weight: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_transpose: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;op_pattern&quot;</span>: <span style=\"color: #008000\">2</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">256</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_transpose&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(p_fc2_weight[v_ax1, v_ax0])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_transpose[v_ax0, v_ax1])\n",
       "                T_transpose[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> p_fc2_weight[v_ax1, v_ax0]\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc1_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), p_fc2_bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            lv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>transpose, (p_fc1_weight,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv_1 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>fused_matmul_add_relu, (x, lv, p_fc1_bias), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv4 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>transpose1, (p_fc2_weight,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>fused_matmul1_add1, (lv_1, lv4, p_fc2_bias), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod = relax.get_pipeline(\"zero\")(mod)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 通用部署IRModule\n",
    "\n",
    "优化完成后，我们可以将模型编译为 TVM 运行时模块。值得注意的是，Apache TVM Unity 提供了通用部署的能力，这意味着可以在不同的后端（包括 CPU、GPU 和其他新兴后端）上部署相同的 IRModule。\n",
    "\n",
    "### 在 CPU 上部署\n",
    "我们可以通过将目标指定为 `llvm` 来在 CPU 上部署 IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.05665974 -0.04947902  0.16557705  0.08111103  0.09358499  0.07472591\n",
      "  -0.03751349 -0.09982338  0.16402456  0.05222905]]\n"
     ]
    }
   ],
   "source": [
    "exec = tvm.compile(mod, target=\"llvm\")\n",
    "dev = tvm.cpu()\n",
    "vm = relax.VirtualMachine(exec, dev)\n",
    "\n",
    "raw_data = np.random.rand(1, 784).astype(\"float32\")\n",
    "data = tvm.runtime.tensor(raw_data, dev)\n",
    "cpu_out = vm[\"main\"](data, *params_from_torch[\"main\"]).numpy()\n",
    "print(cpu_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 在 GPU 上部署\n",
    "\n",
    "除了 CPU 后端，还可以在其他后端上部署 IRModule。例如，可以将 IRModule 部署在 GPU 上。GPU 需要包含额外信息的程序，如线程绑定和共享内存分配。需要进一步的转换来生成 GPU 程序。\n",
    "\n",
    "使用 DLight 来生成 GPU 程序。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm import dlight as dl\n",
    "\n",
    "with tvm.target.Target(\"cuda\"):\n",
    "    gpu_mod = dl.ApplyDefaultSchedule(\n",
    "        dl.gpu.Matmul(),\n",
    "        dl.gpu.Fallback(),\n",
    "    )(mod)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在可以像在 CPU 上那样，在 GPU 上编译 IRModule。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.0566597  -0.04947904  0.16557696  0.08111105  0.09358505  0.07472588\n",
      "  -0.03751348 -0.09982348  0.16402464  0.05222905]]\n"
     ]
    }
   ],
   "source": [
    "exec = tvm.compile(gpu_mod, target=\"cuda\")\n",
    "dev = tvm.device(\"cuda\", 0)\n",
    "vm = relax.VirtualMachine(exec, dev)\n",
    "# Need to allocate data and params on GPU device\n",
    "data = tvm.runtime.tensor(raw_data, dev)\n",
    "gpu_params = [tvm.runtime.tensor(p, dev) for p in params_from_torch[\"main\"]]\n",
    "gpu_out = vm[\"main\"](data, *gpu_params).numpy()\n",
    "print(gpu_out)\n",
    "\n",
    "# Check the correctness of the results\n",
    "assert np.allclose(cpu_out, gpu_out, atol=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 在其他后端上部署\n",
    "\n",
    "Apache TVM Unity 还支持其他后端，如各种类型的 GPU（Metal、ROCm、Vulkan 和 OpenCL）、各种类型的 CPU（x86、ARM）以及其他新兴后端（例如 WebAssembly）。部署过程与 GPU 后端类似。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
